"""
Message Generator模块 - 修复版
功能：生成0/1消息并添加冗余数据
这是FakeTagger系统的第1个模块
"""

import torch
import torch.nn as nn
import numpy as np
import random

class MessageGenerator:
    """
    消息生成器 - 修复版
    负责生成随机的0/1消息，并为消息添加冗余数据以提高鲁棒性
    """
    
    def __init__(self, message_length=30, redundancy_factor=2):
        """
        初始化消息生成器
        
        Args:
            message_length: 原始消息长度（默认30位，与原项目保持一致）
            redundancy_factor: 冗余因子，冗余消息长度 = message_length * redundancy_factor
        """
        self.message_length = message_length
        self.redundancy_factor = redundancy_factor
        self.redundant_length = message_length * redundancy_factor
        
        print(f"信息生成初始化  Message Generator initialized:")
        print(f"  - 源信息长度:     {self.message_length}")
        print(f"  - 冗余因子:       {self.redundancy_factor}")
        print(f"  - 冗余消息长度:   {self.redundant_length}")
        
    def generate_random_message(self, batch_size=1, device='cuda'):
        """
        生成随机的0/1消息
        修复：生成连续值而不是离散值，便于网络训练
        
        Args:
            batch_size: 批大小
            device: 设备类型
            
        Returns:
            torch.Tensor: 原始消息 (batch_size, message_length)，值为接近0或1的连续值
        """
        if isinstance(device, str):
            device = torch.device(device if torch.cuda.is_available() else 'cpu')
        
        # 生成二值化消息（0或1）
        binary_messages = torch.randint(0, 2, (batch_size, self.message_length), 
                                       dtype=torch.float32, device=device)
        
        # 转换为连续值：0 -> 0.1, 1 -> 0.9，避免完全的0/1
        # 这样网络更容易学习，同时保持二值化特性
        continuous_messages = binary_messages * 0.8 + 0.1  # 0->0.1, 1->0.9
        
        # 添加微小的随机扰动，增加训练稳定性
        noise = torch.randn_like(continuous_messages) * 0.01
        continuous_messages = torch.clamp(continuous_messages + noise, 0.0, 1.0)
        
        return continuous_messages
    
    def add_redundancy(self, messages):
        """
        为消息添加冗余数据
        修复：使用更鲁棒的冗余编码策略
        
        Args:
            messages: 原始消息 (batch_size, message_length)
        Returns:
            torch.Tensor: 冗余消息 (batch_size, redundant_length)
        """
        batch_size = messages.shape[0]
        device = messages.device
        
        # 方法1：重复编码（基础冗余）
        repeated_messages = messages.repeat(1, self.redundancy_factor)
        
        # 方法2：添加奇偶校验位（增强鲁棒性）
        if self.redundancy_factor > 1:
            # 计算奇偶校验位
            parity_bits = torch.sum(messages, dim=1, keepdim=True) % 2
            
            # 创建校验增强的冗余消息
            redundant_messages = []
            for i in range(self.redundancy_factor):
                if i == 0:
                    # 第一份：原始消息
                    redundant_messages.append(messages)
                else:
                    # 后续份：原始消息 + 奇偶校验位的组合
                    enhanced_msg = messages.clone()
                    # 在每个奇偶校验位位置添加校验信息
                    if i == 1 and self.message_length > 1:
                        # 第二份：最后一位替换为奇偶校验位
                        enhanced_msg[:, -1] = parity_bits.squeeze()
                    redundant_messages.append(enhanced_msg)
            
            redundant_messages = torch.cat(redundant_messages, dim=1)
        else:
            redundant_messages = repeated_messages
        
        # 添加适量的噪声（模拟传输过程中的扰动）
        noise_scale = 0.02  # 增加噪声强度
        noise = torch.randn_like(redundant_messages) * noise_scale
        redundant_messages = torch.clamp(redundant_messages + noise, 0.0, 1.0)
        
        return redundant_messages
    
    def generate_redundant_message(self, batch_size=1, device='cuda'):
        """
        一步生成带冗余的消息
        
        Args:
            batch_size: 批大小
            device: 设备类型
            
        Returns:
            tuple: (原始消息, 冗余消息)
        """
        original_messages = self.generate_random_message(batch_size, device)
        redundant_messages = self.add_redundancy(original_messages)
        
        return original_messages, redundant_messages

def test_message_generator():
    """测试消息生成器功能"""
    print("=== Testing Message Generator ===")
    
    # 创建生成器（使用项目默认的message_size=30）
    generator = MessageGenerator(message_length=30, redundancy_factor=2)
    
    # 测试生成消息
    batch_size = 4
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"设备: {device}")

    # 生成原始消息
    original_msgs = generator.generate_random_message(batch_size, device)
    print(f"原始信息 形状shape:  {original_msgs.shape}")
    print(f"原始信息 样本sample: {original_msgs[0][:10].cpu().numpy()}")
    print(f"原始信息 范围range: [{original_msgs.min().item():.3f}, {original_msgs.max().item():.3f}]")
    
    # 添加冗余
    redundant_msgs = generator.add_redundancy(original_msgs)
    print(f"冗余信息 形状shape: {redundant_msgs.shape}")
    print(f"冗余信息 样本sample(前10bit): {redundant_msgs[0][:10].cpu().numpy()}")
    print(f"冗余信息 范围range: [{redundant_msgs.min().item():.3f}, {redundant_msgs.max().item():.3f}]")
    
    # 一步生成
    orig, redun = generator.generate_redundant_message(batch_size, device)
    print(f"一步生成 - 原始: {orig.shape}, 冗余: {redun.shape}")
    
    # 测试二值化恢复
    binary_orig = (orig > 0.5).float()
    print(f"二值化恢复 - 特征值: {torch.unique(binary_orig)}")
    
    print("Test completed!\n")

if __name__ == "__main__":
    test_message_generator() 